"""The model on the node."""

import SocketServer
import json
import logging
import os
import sys
import traceback
from gaweb import consts
from gaweb.node import excepts
from gaweb.node import helpers

import warnings
warnings.filterwarnings('ignore')
import grads
warnings.resetwarnings()


class NodeServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
    """A Node Server."""
    def __init__(self):
        if os.environ.get('USER') == 'root':
            raise excepts.NodeServerError('Refusing to run as root')
        config = helpers.read_config()
        SocketServer.TCPServer.__init__(self, (config['address'], consts.NODE_LISTEN_PORT), RequestHandler)


class RequestHandler(SocketServer.BaseRequestHandler):
    """Request Handler, instantiated once for each connection.

    In the master/node model, one connection (and therefore one of these) is
    created for each active session.
    """
        
    def _get_dims(self):
        """Returns the current dims situation."""
        try:
            qh = self.ga.query('dims')
            dims = helpers.handle_to_dict(qh, filter=True)
        except grads.GrADSError:
            dims = None
        
        return dims

    def _get_ctlinfo(self):
        """Returns the current ctlinfo situation."""
        try:
            qh = self.ga.query('ctlinfo')
            ctlinfo = helpers.handle_to_dict(qh, filter=True)

            # Let's inject a little bit extra. Our convention is to put a '_' in front
            # of our keys.
            dims = self._get_dims()
            old_t = dims['t']
            self.ga.cmd('set t %d' % ctlinfo['nt'])

            new_dims = self._get_dims()
            ctlinfo['_tf'] = new_dims['time'][0]
            self.ga.cmd('set t %d %d' % old_t)

        except grads.GrADSError:
            ctlinfo = None
        
        return ctlinfo

    def handle(self):
        self.session_id = self.request.recv(consts.LEN_SESSION_ID)
        logging.info('New session from %s (session_id=%s)' % (self.client_address[0], self.session_id))

        self.ga = grads.GrADS(Echo=False, Window=False)

        self.ga.cmd('set gxout shaded')

        # Loop while true, dispatching messages to instance methods that return a
        # dictionary of data for the message response. Each request is blocking.
        buff = ''
        while True:
            buff += self.request.recv(1024)
            if not buff:
                # We handle the master disconnecting by returning from the handle method
                # (which ends the connection on our side).                
                logging.info('Master closed connection. Ending session in node.')
                return

            logging.debug('%s\t<-- RECV: %s' % (self.client_address[0], repr(buff)))
            msgs = buff.split('\n')
            buff = msgs.pop(-1)

            for msg in msgs:
                try:
                    msg_type, data = msg.split('::', 1)
                except ValueError:
                    logging.error('Recieved unsplittable message: %s' % repr(msg))
                    continue
                try:
                    data = json.loads(data)
                except ValueError:
                    logging.error('Could not parse JSON data: %s' % repr(data))
                    continue

                method = getattr(self, 'handle_%s' % msg_type, None)
                if method:
                    try:
                        resp_data = method(data)
                    except Exception, e:
                        resp_data = {'error': -1}
                        exc_type, exc_value, exc_traceback = sys.exc_info()
                        logging.error(traceback.format_exc(exc_traceback))
                        logging.error(repr(e))
                    self.send(msg_type, resp_data)
                else:
                    logging.error('No method found for message type %s' % repr(msg_type))

    def send(self, msg_type, data):
        msg = '%s::%s\n' % (msg_type, json.dumps(data))
        logging.debug('%s\t--> SEND: %s' % (self.client_address[0], repr(msg)))
        self.request.send(msg)

    def handle_command(self, data):
        """Handle a command message.

        The tricky thing to know here is that the variable 'command' is actually a
        newline-separated sequence of GrADS statements.
        """
        command = data['command']

        # Collect list of first tokens in sequence of commands. This is useful for
        # later analysis.
        first_toks = [line.strip().split(' ')[0] for line in command.split('\n') if line]
        first_toks = [tok.lower() for tok in first_toks]

        already_found_display_expr = None

        # We intervene if:
        # 1) The command sequence is too many characters long
        # 2) The command sequence contains a block command (we do a full reject in this case)
        if len(command) > consts.MAX_COMMAND_LEN:
            output = 'Command to long (>%d characters)' % consts.MAX_COMMAND_LEN
        elif any(first_tok in consts.BLOCKED_COMMANDS for first_tok in first_toks):
            output = 'You are not permitted to use this command.'
        else:
            # Execute the command
            logging.info('Session %s; Running command %s' % (self.session_id, repr(command)))

            try:
                self.ga.cmd(command)
                n = self.ga.rline()
                output = '\n'.join(self.ga.rline(i+1) for i in range(n))                
            except grads.GrADSError:
                n = self.ga.rline()
                output = '\n'.join(self.ga.rline(i+1) for i in range(n))
        
        # Build response data based on analysis of commands
        response_data = {'output': output}

        if any(first_tok in consts.DISPLAY_COMMANDS for first_tok in first_toks):
            response_data['display_expr_changed'] = True

            if already_found_display_expr:
                response_data['display_expr'] = already_found_display_expr
            else:
                for line in reversed(command.split('\n')):
                    toks = line.split(' ')
                    if len(toks) >= 2:
                        first = toks[0]
                        rest = ' '.join(toks[1:])
                    
                        if first.lower() in consts.DISPLAY_COMMANDS:
                            response_data['display_expr']  = rest
                            break
                if 'display_expr' not in response_data:            # Caused by bad command
                    response_data['display_expr'] = '' 

        else:
            response_data['display_expr_changed'] = False
            response_data['display_expr'] = None

        if any(first_tok in consts.REFRESH_COMMANDS for first_tok in first_toks):
            dims = self._get_dims()
            if 1 < dims['nt'] < consts.MAX_ANIMATION_LEN and response_data['display_expr_changed']:
                # Redo our display request as an animation
                img_fnames = []
                command = ''
                for t in range(dims['t'][0], dims['t'][1] + 1):
                    img_fname = '/tmp/%s-%d.png' % (self.session_id, t)
                    img_fnames.append(img_fname)
                    command += 'set t %d\n' % t
                    command += 'clear\n'
                    command += 'd %s\n' % response_data['display_expr']
                    command += 'draw title t = %d\n' % t
                    command += 'cbarn\n'
                    command += 'printim %s x588 y376\n' % img_fname
                command += 'set t %d %d\n' % dims['t']
                self.ga.cmd(command)

                anim_fname = '/tmp/%s-anim.gif' % self.session_id
                os.system('convert -delay 20 -loop 0 %s %s' % (' '.join(img_fnames), anim_fname))
                
                base64_data = open(anim_fname).read().encode('base64')
                response_data['image_data'] = 'data:image/gif;base64,%s' % base64_data
                response_data['image'] = True
            else:
                # Normal image output
                img_fname = '/tmp/%s.png' % self.session_id
                self.ga.cmd('printim %s x588 y376' % img_fname)
                base64_data = open(img_fname).read().encode('base64')
                response_data['image_data'] = 'data:image/gif;base64,%s' % base64_data
                response_data['image'] = True
        else:
            response_data['image'] = False
            response_data['image_data'] = None

        if any(first_tok in consts.REFRESH_FILES_COMMANDS for first_tok in first_toks):
            response_data['files_changed'] = True
        else:
            response_data['files_changed'] = False

        if any(first_tok in consts.REFRESH_DIMS_COMMANDS for first_tok in first_toks):
            response_data['dims_changed'] = True
            response_data['dims'] = self._get_dims()
            response_data['ctlinfo'] = self._get_ctlinfo()
        else:
            response_data['dims_changed'] = False
            response_data['dims'] = None
            response_data['ctlinfo'] = None
        
        return response_data

    def handle_diminfo(self, data):
        """Handles a diminfo request; this is essentially a request
        to ga.query('dims').
        
        Queries are only successful if one or more files are open.

        If the query is successful, returns a dictionary {dims: dims_dict, ctlinfo: ctlinfo_dict}.
        If the query is unsuccessful, returns a dictionary {dims: None, ctlinfo: None}.
        """            
        response_data = {'dims': self._get_dims(),
                         'ctlinfo': self._get_ctlinfo()}
        return response_data

    def handle_nodestats(self, data):
        """Handle a nodestats request from the master."""
        response_data = {}
        response_data['text'] = helpers.get_system_info_text()
        
        return response_data

    def handle_varinfo(self, data):
        """Handles a varinfo request; this is information about all variables
        present in all open files.
        
        Returns an array of size 5. If array[i] is non-null, then file #(i + 1) is
        open. In this case, array[i] is an array of 4-tuples.
        """
        files = []

        for i in range(1, 6):
            try:
                handle = self.ga.query('file %d' % i)
            except grads.GrADSError:
                handle = None
            
            if handle is None:
                files.append(None)
            else:
                files.append(handle.var_info)
        
        response_data = {}
        response_data['files'] = files

        return response_data
